import time, json

from openai import OpenAI
from config.config import OPENAPI_API_KEY_4, OPENAPI_API_BASE_4

client = OpenAI(
    api_key=OPENAPI_API_KEY_4,
    base_url=OPENAPI_API_BASE_4
)


## 故事agent 执行流程
## 1、定义一个生成故事的 function
## 2、定义一个故事格式化 function
## 3、定一个发布到博客的 fcuntion
def generateStory(theme: str):
    return f"根据你的主题{theme},我们生成了这样一个故事：小白兔白又白，两只耳朵竖起来，爱吃萝卜爱吃菜，蹦蹦跳跳真可爱"


def formatStory(content: str):
    return f"小白兔白又白，两只耳朵竖起来，爱吃萝卜爱吃菜，蹦蹦跳跳真可爱"


def publishStory(content: str):
    return f"我们已经将故事发布到博客"


tools = [
    {
        "type": "function",
        "function": {
            "name": "generateStory",
            "description": "这是一个根据传入指定主题生成故事的方法",
            "parameters": {
                "type": "object",
                "properties": {
                    "theme": {
                        "type": "string",
                        "description": "故事具体主题"
                    }
                },
                "required": ["theme"]
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "formatStory",
            "description": "格式化传入的故事内容为统一格式，并返回格式化后的故事",
            "parameters": {
                "type": "object",
                "properties": {
                    "content": {
                        "type": "string",
                        "description": "生成的故事的详细内容"
                    }
                },
                "required": ["content"]
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "publishStory",
            "description": "拿到格式化好的故事发布到博客",
            "parameters": {
                "type": "object",
                "properties": {
                    "content": {
                        "type": "string",
                        "description": "格式化好的具体故事内容"
                    }
                },
                "required": ["content"]
            }
        }
    },
]
functionMaps = {
    "generateStory": generateStory,
    "formatStory": formatStory,
    "publishStory": publishStory,
}


def createStory(theme: str = "小兔子找朋友"):
    messages = [
        {
            "role": "user", "content": f"主题是：{theme},生成一篇故事，并发布到博客"
        }
    ]
    while True:
        response = client.chat.completions.create(
            model="gpt-4o-mini",
            messages=messages,
            tools=tools,
            tool_choice="auto"
        )
        if response.choices[0].message.tool_calls:
            for tool in response.choices[0].message.tool_calls:
                args = json.loads(tool.function.arguments)
                result = functionMaps[tool.function.name](**args)
                messages.append({
                    "role": "assistant",
                    "tool_calls": [tool]
                })
                messages.append({
                    "role": "tool",
                    "name": tool.function.name,
                    "tool_call_id": tool.id,
                    "content": result
                })
                time.sleep(1)
                print(result)
        else:
            print("无任何工具可以触发，结束-----------")
            print(response.choices[0].message.content)
            break
    return {"content": response.choices[0].message.content}
